'''
File: parsingGenbankFiles.py
Author: Adam Hockenberry
Description: 
'''

#Standard path imports
from optparse import OptionParser
import sys

#Non-standard imports
from Bio import SeqIO
from sequenceAnalysisFunctions import *

def load_defined_dataset(definedDatasetFile):
#if you want to only look at a subset of genes
#file layout is just one gene name per line with no puncuation, tabs, etc
#code only currently set-up to deal with gene names in the form of locus_tag 
    f = open(definedDatasetFile).readlines() 
    for x in range(len(f)):
        f[x] = f[x].strip('\n')
        f[x] = f[x].strip() #strip leading/trailing spaces if any are there
    return f #where f is a list of strings


def has_internal_stop_codon(sequenceString):
#make sure that TAA, TAG, TGA don't appear more than once in the string
    sequenceAsList = split_string_into_list_of_triplets(sequenceString)
    if sequenceAsList.count('TAA') + sequenceAsList.count('TAG') + sequenceAsList.count('TGA') > 1:
        return True
    return False


def contains_strange_base(sequenceString):
#make sure there is no bases or markups other than ATGC
    totalCounts = len(sequenceString)
    countsOfA = sequenceString.count('A')
    countsOfT = sequenceString.count('T')
    countsOfG = sequenceString.count('G')
    countsOfC = sequenceString.count('C')
    if countsOfA + countsOfT + countsOfG + countsOfC != totalCounts:
        return True
    return False


def is_protein_coding_sequence(sequenceString):
#tests to make sure the sequence is legitimate with no funny business
    testResult = True
    if sequenceString[0:3] not in ['ATG', 'GTG', 'TTG'] or sequenceString[-3:] not in ['TAA', 'TAG', 'TGA']:
        print "The given coding sequence does not have either a start or stop codon:"
        testResult = False
        return testResult
    if len(sequenceString) % 3. != 0:
        print "The given coding sequence is not a multiple of 3 in length:"
        testResult = False
        return testResult
    if has_internal_stop_codon(sequenceString):
        print "The given coding sequence contains internal stop codons:"
        testResult = False
        return testResult
    if contains_strange_base(sequenceString):
        print "The given coding sequence contains a letter other than ATGC:"
        testResult = False
        return testResult
    return testResult

def check_for_redundancy(sequenceList, sequenceNames):
#make sure a gene isn't represented twice
#a lot are, not sure whether to care or deal with this 
#if they are legitimate gene duplications which they appear to be
    listOfRedundantIndices = []
    for x in range(len(sequenceList)):
        for y in range(x+1,len(sequenceList),1):
            if sequenceList[x] == sequenceList[y]:
                print 'Error: ', sequenceNames[x], '(', x, ')is the same as ', sequenceNames[y], '(', y, ')'
    return sequenceList 



####################################################################################
#The following code works with single genbank file genomes. It must be run separately for each chromosome in
#multiple chromosomal organisms and the resulting lists concatenated
#will write code to do that eventually
####################################################################################
def get_coding_sequences(genbankFileLocation, subsetFileLocation=False, worryAboutRedundancy=False):
    sequenceList = []
    lociNames = []
    errorLoci = []
    if subsetFileLocation:
        subsetNames = load_defined_dataset(subsetFileLocation)
    else:
        subsetNames = []
####################################################################################
#reading and iterating through the genbank file
    genome = SeqIO.read(genbankFileLocation, 'genbank')
    for feature in genome.features:
        if feature.type == 'CDS':
            locus = feature.qualifiers['locus_tag'][0]
            if not subsetFileLocation:
                subsetNames.append(locus)
            if locus in subsetNames: 
            #note my subset file references only the locus_tag for genes of interest
                seq = str(feature.extract(genome.seq))
                if is_protein_coding_sequence(seq):
                    sequenceList.append(seq)
                    lociNames.append(locus)
                else:
                    print locus
####################################################################################
#some sanity checks
    if subsetFileLocation:
        if len(sequenceList) != len(subsetNames):
            print "Number of extracted sequences does not match number expected in subset"
    print "##################################################################################"
    print "total number of sequences extracted from genome is:", len(sequenceList)
    if worryAboutRedundancy:
        sequenceList = check_for_redundancy(sequenceList, lociNames)
    return sequenceList

if __name__ == '__main__':
    genomeFile = sys.argv[1]
    if len(sys.argv) > 2:
        subsetFile = sys.argv[2]
    else:
        subsetFile = False
    get_coding_sequences(genomeFile, subsetFile)
